{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZCVp1FbaJEH7"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bm-P_UOxJQQ8"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ppNrt1-RK8SO"
      },
      "source": [
        "# Train embeddings on TPU using Autoencoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8L8kFn5ILCGv"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This colab explores how to train autoencoders on a TPU device.\n",
        "\n",
        "For this colab, consider the following scenario: you have an image classification model that you want to improve by adding some additional features. The features that you can add to the model could be image embeddings that can be separately trained on a TPU.\n",
        "\n",
        "This example uses a fully-connected one layer model as the model that you want to make better with additional features trained on a TPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i5HlRfWwLMGZ"
      },
      "source": [
        "## Learning objectives\n",
        "\n",
        "In this Colab, you will learn how to\n",
        "* Build a fully-connected one layer model to classify images\n",
        "* Build an autoencoder and train on those images, in an unsupervised fashion, to produce image embeddings\n",
        "* Retrain a fully-connected one layer model with additonal features, the embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Boy46RrCZtTp"
      },
      "source": [
        "### Check that you have enabled TPUs in this notebook\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xx7NDPvMSPMF"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit \u003e Notebook settings \u003e Hardware accelerator'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JeMli385le2A"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from absl import logging\n",
        "\n",
        "logging.set_verbosity(logging.ERROR)\n",
        "\n",
        "# Initialize TPU Strategy.\n",
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "T3abQXrfZ94y"
      },
      "source": [
        "### Get data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "EB2PiqKS9xMh"
      },
      "outputs": [],
      "source": [
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "y_train, y_test = y_train.astype(np.int32), y_test.astype(np.int32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7MbVCS56aBjJ"
      },
      "source": [
        "### Function to visualize our images and pick the first image from the test set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6iNexrDZVFR9"
      },
      "outputs": [],
      "source": [
        "def show_img(img):\n",
        "  plt.figure()\n",
        "  plt.imshow(img)\n",
        "  plt.grid(False)\n",
        "  plt.show()\n",
        "\n",
        "img = 0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "K1YIN0l3aVCS"
      },
      "source": [
        "### The first image from the test set is the number 7"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vKFsymlkALr8"
      },
      "outputs": [],
      "source": [
        "show_img(x_test[img].reshape(28, 28))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YSrG-xNbaa_C"
      },
      "source": [
        "### MNIST setup\n",
        "There are 10 classes (one for each digit) and each image is 28 by 28 pixels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "LF0CFIcfGH4z"
      },
      "outputs": [],
      "source": [
        "NUM_CLASSES = 10\n",
        "\n",
        "# input image dimensions\n",
        "IMG_ROWS, IMG_COLS = 28, 28\n",
        "\n",
        "x_train = x_train.reshape(x_train.shape[0], IMG_ROWS, IMG_COLS, 1)\n",
        "x_test = x_test.reshape(x_test.shape[0], IMG_ROWS, IMG_COLS, 1)\n",
        "\n",
        "x_train = x_train.astype('float32')\n",
        "x_test = x_test.astype('float32')\n",
        "x_train = x_train / 255.0\n",
        "x_test = x_test / 255.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LulbuII9lqng"
      },
      "source": [
        "### Original model\n",
        "\n",
        "Here is a contrived example where the training happens only on the corners of the MNIST image.\n",
        "\n",
        "Suppose that your original model, the fully-connected one layer network, was too computationally heavy, in terms of resources, and thus you could only afford to train on parts of the images. Instead of training on 28 by 28 pixels (784 pixels), you train on 14 by 14 pixels (196 pixels). This colab will later show that just by adding 49 more pixels to each training example, the size of each embedding, accuracy can be significantly increased.\n",
        "\n",
        "This way you introduce minimal changes to an original model while gaining benefits from a heavy computational task that you can be offload to a TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "m5E6fMzviiDB"
      },
      "outputs": [],
      "source": [
        "x_train_corners = x_train[:, :14, :14, :]\n",
        "x_test_corners = x_test[:, :14, :14, :]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mggEYZbxh0Tx"
      },
      "source": [
        "### The first image corner from the test set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sdYcDDZQiBOP"
      },
      "outputs": [],
      "source": [
        "show_img(x_test_corners[img].reshape(14, 14))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y-4XhCQ7bO39"
      },
      "source": [
        "### Create a model with one fully-connected layer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6hBRXkNyCGpV"
      },
      "outputs": [],
      "source": [
        "def get_model(input_shape):\n",
        "  ip = tf.keras.layers.Input(shape=input_shape)\n",
        "  x = tf.keras.layers.Flatten()(ip)\n",
        "  x = tf.keras.layers.Dense(NUM_CLASSES, activation='sigmoid')(x)\n",
        "  \n",
        "  model = tf.keras.models.Model(ip, x)\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6DXBZSVaJ15R"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model0 = get_model(x_train_corners[0].shape)\n",
        "  model0.compile(\n",
        "      optimizer=tf.optimizers.SGD(learning_rate=0.05), \n",
        "      loss=tf.losses.SparseCategoricalCrossentropy(), \n",
        "      metrics=[tf.metrics.SparseCategoricalAccuracy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ve6I4U8ybaun"
      },
      "source": [
        "### Train and evaluate the fully-connected one layer model on CPU\n",
        "\n",
        "As expected, the model performs poorly, but it does train fairly quickly. Expected accuracy is 65%."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "So-1S36cKs1u"
      },
      "outputs": [],
      "source": [
        "model0.fit(x_train_corners, y_train, epochs=3, batch_size=128)\n",
        "model0.evaluate(x_test_corners, y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vY3mmuZCb3BD"
      },
      "source": [
        "### Create an autoencoder and make sure to get back an encoder as well"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WwqpI0kMIb9V"
      },
      "outputs": [],
      "source": [
        "def get_autoencoder_and_encoder(input_shape):\n",
        "  ip = tf.keras.layers.Input(shape=input_shape)\n",
        "\n",
        "  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(ip)\n",
        "  x = tf.keras.layers.MaxPooling2D((2, 2), padding='same')(x)\n",
        "  x = tf.keras.layers.Conv2D(1, (3, 3), activation='relu', padding='same')(x)\n",
        "  encoded = tf.keras.layers.MaxPooling2D((2, 2), padding='same')(x)\n",
        "\n",
        "  x = tf.keras.layers.Conv2DTranspose(1, (3, 3), activation='relu', strides=2, padding='same')(encoded)\n",
        "  x = tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', strides=2, padding='same')(x)\n",
        "  \n",
        "  decoded = tf.keras.layers.Conv2DTranspose(1, (3, 3), activation='sigmoid', padding='same')(x)\n",
        "\n",
        "  autoencoder = tf.keras.models.Model(ip, outputs=decoded)\n",
        "  encoder = tf.keras.models.Model(ip, encoded)\n",
        "  \n",
        "  return autoencoder, encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oGNnC4FwcDoX"
      },
      "source": [
        "### Train the autoencoder on TPU\n",
        "\n",
        "This is a computationally resource expensive operation that can be offloaded to the TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "j-2PBN7smmCy"
      },
      "outputs": [],
      "source": [
        "tf.keras.backend.clear_session()\n",
        "\n",
        "with strategy.scope():\n",
        "  autoencoder, encoder = get_autoencoder_and_encoder(x_train[0].shape)\n",
        "  \n",
        "  autoencoder.compile(\n",
        "    optimizer=tf.optimizers.Adam(), \n",
        "    loss=tf.keras.losses.BinaryCrossentropy(),\n",
        "    metrics=[tf.metrics.BinaryAccuracy()])\n",
        "  \n",
        "  autoencoder.fit(\n",
        "    x_train,\n",
        "    x_train, \n",
        "    batch_size=128,\n",
        "    epochs=3,\n",
        "    steps_per_epoch=468,\n",
        "    validation_data=(x_test, x_test))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CNbsGydbcLhd"
      },
      "source": [
        "### Produce image embeddings\n",
        "\n",
        "Now that the autoencoder is trained, you can use the encoder part to produce image embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VQ_mrpWPV86h"
      },
      "outputs": [],
      "source": [
        "x_train_embeddings = encoder.predict(x_train)\n",
        "x_test_embeddings = encoder.predict(x_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qzNASLGncnbV"
      },
      "source": [
        "### Produce image reconstructions\n",
        "\n",
        "Let's visually see the quality of our autoencoder to see how the number 7 from above is reconstructed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "44Dqs3n-E5RL"
      },
      "outputs": [],
      "source": [
        "x_test_hat = autoencoder.predict(x_test[:8])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Y3HBZLGDdGcQ"
      },
      "source": [
        "### Reconstructed number 7\n",
        "\n",
        "This looks like the number 7 so now you can be confident in the quality of our embeddings. The autoencoder learned to compress and uncompress information accurately."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "X2_UhQQ3yTZJ"
      },
      "outputs": [],
      "source": [
        "show_img(x_test_hat[img].reshape(28, 28))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "WRBZDl0Ddjbd"
      },
      "source": [
        "### Check the original image\n",
        "\n",
        "Remember, the image in the previous section is the reconstructed image. Compare it to the original image, as shown here.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uoJctld6yekJ"
      },
      "outputs": [],
      "source": [
        "show_img(x_test[0].reshape(28, 28))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Y3pjFtJokLt4"
      },
      "source": [
        "### Examine the embedding for the number 7"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "k1QsGM7okUC4"
      },
      "outputs": [],
      "source": [
        "show_img(x_test_embeddings[0].reshape(7, 7))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AGV71mzWkajN"
      },
      "source": [
        "### Augment the corners dataset\n",
        "\n",
        "The following code augments the corners dataset with embeddings trained on TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0Iby-JLwECuf"
      },
      "outputs": [],
      "source": [
        "x_train_augmented = np.concatenate([x_train_corners.reshape(60000, 14*14, 1), x_train_embeddings.reshape(60000, 7*7, 1)], axis=1)\n",
        "x_test_augmented = np.concatenate([x_test_corners.reshape(10000, 14*14, 1), x_test_embeddings.reshape(10000, 7*7, 1)], axis=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TFnK0ADuklZy"
      },
      "source": [
        "### Retrain the original model\n",
        "\n",
        "At this point, you can train the original model using the augmented dataset. You should verify that the TPU embeddings augmented model works better than without embeddings. Expected accuracy is 87%."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Di-ub4eKEYdf"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model1 = get_model(x_train_augmented[0].shape)\n",
        "  model1.compile(\n",
        "      optimizer=tf.optimizers.SGD(learning_rate=0.06), \n",
        "      loss=tf.losses.SparseCategoricalCrossentropy(), \n",
        "      metrics=['accuracy'])\n",
        "\n",
        "model1.fit(x_train_augmented, y_train, epochs=3, batch_size=128)\n",
        "model1.evaluate(x_test_augmented, y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "OTjeGkOFsH0Q"
      },
      "source": [
        "## What's next\n",
        "\n",
        "* Learn about [Cloud TPUs](https://cloud.google.com/tpu/docs) that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.\n",
        "* Explore the range of [Cloud TPU tutorials and Colabs](https://cloud.google.com/tpu/docs/tutorials) to find other examples that can be used when implementing your ML project.\n",
        "\n",
        "On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured [deep learning VMs](https://cloud.google.com/deep-learning-vm/),  you will find [AutoML](https://cloud.google.com/automl/)*(beta)* for training custom models without writing code and [Cloud ML Engine](https://cloud.google.com/ml-engine/docs/) which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware.\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Train embeddings on TPU using Autoencoder",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
